#!/usr/bin/env python
"""
GitCode MCP客户端测试脚本
"""
import asyncio
import json
import os
import sys
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client

async def test_mcp_client():
    """测试MCP客户端连接和API调用"""
    print("╔════════════════════════════════════════════════════════╗")
    print("║        GitCode MCP 客户端测试                          ║")
    print("╚════════════════════════════════════════════════════════╝")
    print()
    
    # 检查环境变量
    gitcode_token = os.environ.get("GITCODE_TOKEN")
    if not gitcode_token:
        print("⚠️  警告: 未设置 GITCODE_TOKEN 环境变量")
        print("   某些测试将被跳过")
        print()
    
    # 配置服务器参数
    server_params = StdioServerParameters(
        command=sys.executable,  # 使用当前 Python 解释器
        args=["-m", "gitcode_mcp"],
        env={"GITCODE_TOKEN": gitcode_token} if gitcode_token else {}
    )
    
    print("🔌 正在连接 GitCode MCP 服务器...")
    
    try:
        async with stdio_client(server_params) as (read, write):
            async with ClientSession(read, write) as session:
                # 初始化会话
                await session.initialize()
                print("✅ 成功连接到服务器！")
                print()
                
                # 1. 列出可用工具
                print("═" * 60)
                print("📋 可用工具")
                print("═" * 60)
                tools_result = await session.list_tools()
                print(f"共有 {len(tools_result.tools)} 个工具:")
                for tool in tools_result.tools:
                    desc = tool.description or "无描述"
                    print(f"  • {tool.name}")
                    print(f"    └─ {desc[:60]}...")
                print()
                
                # 2. 列出可用资源
                print("═" * 60)
                print("📦 可用资源")
                print("═" * 60)
                resources_result = await session.list_resources()
                if resources_result.resources:
                    print(f"共有 {len(resources_result.resources)} 个资源:")
                    for resource in resources_result.resources:
                        print(f"  • {resource.uri}")
                        if hasattr(resource, 'name') and resource.name:
                            print(f"    └─ {resource.name}")
                else:
                    print("  (无可用资源)")
                print()
                
                # 3. 列出可用提示
                print("═" * 60)
                print("💬 可用提示模板")
                print("═" * 60)
                prompts_result = await session.list_prompts()
                if prompts_result.prompts:
                    print(f"共有 {len(prompts_result.prompts)} 个提示模板:")
                    for prompt in prompts_result.prompts:
                        desc = prompt.description or "无描述"
                        print(f"  • {prompt.name}")
                        print(f"    └─ {desc[:60]}...")
                else:
                    print("  (无可用提示模板)")
                print()
                
                # 4. 测试 API 调用（仅在有 token 时）
                if gitcode_token:
                    print("═" * 60)
                    print("🧪 测试 API 调用")
                    print("═" * 60)
                    
                    try:
                        # 测试列出仓库
                        print("测试: list_repositories")
                        result = await session.call_tool("list_repositories", arguments={})
                        repos = json.loads(result.content[0].text)
                        print(f"✅ 成功获取 {len(repos)} 个仓库")
                        
                        if repos:
                            print("\n前 5 个仓库:")
                            for repo in repos[:5]:
                                full_name = repo.get('full_name', 'unknown')
                                desc = repo.get('description', '无描述')
                                print(f"  • {full_name}")
                                if desc:
                                    print(f"    └─ {desc[:50]}")
                            
                            if len(repos) > 5:
                                print(f"\n  ... 及其他 {len(repos) - 5} 个仓库")
                            
                            # 测试获取单个仓库详情
                            if repos:
                                print("\n测试: get_repository")
                                first_repo = repos[0]
                                owner, repo_name = first_repo['full_name'].split('/')
                                result = await session.call_tool(
                                    "get_repository",
                                    arguments={"owner": owner, "repo": repo_name}
                                )
                                repo_data = json.loads(result.content[0].text)
                                print(f"✅ 成功获取仓库详情: {repo_data.get('full_name')}")
                                print(f"   星标数: {repo_data.get('stargazers_count', 0)}")
                                print(f"   Fork数: {repo_data.get('forks_count', 0)}")
                        
                    except Exception as e:
                        print(f"❌ API 调用失败: {e}")
                        import traceback
                        traceback.print_exc()
                    
                    print()
                else:
                    print("═" * 60)
                    print("⚠️  跳过 API 测试（未设置 GITCODE_TOKEN）")
                    print("═" * 60)
                    print()
                
    except Exception as e:
        print(f"\n❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False
    
    print("═" * 60)
    print("✅ 所有测试通过！")
    print("═" * 60)
    return True

if __name__ == "__main__":
    success = asyncio.run(test_mcp_client())
    sys.exit(0 if success else 1) 